import numpy as np
from collections import defaultdict
from copy import deepcopy
import matplotlib.pyplot as plt


def kmeans(centers, sets):
    flag = True
    condition = None
    bracket = defaultdict(list)
    centers = list(centers)
    while flag:
        bracket.clear()
        for point in sets:
            dists = np.array([np.linalg.norm(center-point, ord=2) for center in centers])
            idx = np.argmin(dists)
            bracket[idx].append(list(point))
        # print(bracket)
        new_center = []
        for key, val in bracket.items():
            tmp = np.array(val)
            new_center.append(tmp.mean(axis=0))
        centers = np.array(new_center)

        if condition == bracket:
            flag = False
        condition = deepcopy(bracket)
    res = []
    for k, v in bracket.items():
        res.append(np.array(v))
    return centers, res


if __name__ == '__main__':

    data = np.array([
            [0.4, 0.3],
            [0, 1],
            [2, 2],
            [0, 2],
            [0.13, 0.34],
            [0.3, 0.6],
            [0.4, 0.67],
            [0.4, 0.92],
            [0.61, 0.82],
            [0.63, 0.62],
            [0.58, 0.72],
            [0.41, 0.52],
            [0.74, 0.83],
            [0.81, 1.21],
            [0.71, 0.3],
            [0.78, 0.32],
            [0.81, 0.1],
            [0.2, 0.52],
            [2.1, 0.7],
            [1.3, 1.4],
            [1.5, 1.2],
            [1.45, 1.31],
            [1.61, 1.52],
            [1.47, 1.47],
            [1.4, 1.1],
            [1.3, 1.2],
            [1.7, 1.5],
            [1.2, 1.8],
            [1.5, 1.5],
            [1.5, 1.8],
            [1, 1],
            [0.7, 0.8],
            [0.3, 2.1],
            [1.2, 1.7],
            [2.2, 2.5],
            [0.7, 1.4]
        ])

    centers = np.array([
        [0., 0.],
        [0.5, 0.5],
        [1.1, 1.3]
    ])
    new_centers, res = kmeans(centers, data)
    print(len(res))
